"""
股票数据爬虫模块

该模块负责从Tushare API获取股票数据并存储到MySQL数据库中。
主要功能包括：
1. 初始化数据库连接
2. 创建股票数据表(如果不存在)
3. 获取上证指数日线数据
4. 将数据存储到MySQL数据库
"""
from datetime import datetime
import pandas as pd
import tushare as ts
from sqlalchemy import create_engine, text
from config import TS_TOKEN, DB_CONFIG
from technical_indicators.indicators import calculate_technical_indicators

# 初始化tushare
pro = ts.pro_api(TS_TOKEN)


def get_db_engine():
    """
    创建并返回数据库引擎对象
    
    返回:
        sqlalchemy.engine.Engine: 数据库引擎实例
    """
    return create_engine(
        f"mysql+mysqlconnector://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"
    )


def main():
    """
    主程序入口
    
    功能:
    1. 初始化数据库连接
    2. 检查并创建股票数据表
    3. 获取上证指数日线数据
    4. 将新数据存储到数据库
    """
    try:
        # 初始化数据库连接
        engine = get_db_engine()
        
        # 检查表是否存在，不存在则创建
        with engine.connect() as conn:
            result = conn.execute(text(
                "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES "
                "WHERE TABLE_NAME = 'stock_index'"
            ))
            if not result.fetchone():
                conn.execute(text(
                    """
                    CREATE TABLE stock_index (
                        id INT AUTO_INCREMENT PRIMARY KEY,
                        ts_code VARCHAR(20),
                        trade_date DATE,
                        open FLOAT,
                        high FLOAT,
                        low FLOAT,
                        close FLOAT,
                        pre_close FLOAT,
                        `change` FLOAT,
                        pct_chg FLOAT,
                        vol FLOAT,
                        amount FLOAT,
                        result1 FLOAT
                    )
                    """
                ))
                print('成功创建stock_index表')

        # 获取上证指数日线数据
        # 获取数据并标准化日期格式
        df = pro.daily(ts_code='300657.SZ', start_date='20220301', end_date=datetime.now().strftime('%Y%m%d'))
        print('获取到的原始数据:')
        print(df)
        
        # 转换日期格式
        df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
        print('转换后的日期格式:')
        print(df['trade_date'])
        
        # 计算技术指标(当前注释掉)
        # df = calculate_technical_indicators(df)

        # 保存到数据库
        # 删除当天已有数据避免重复
        existing_dates = pd.read_sql('SELECT trade_date FROM stock_index', engine)
        print('数据库中已有日期数据:')
        print(existing_dates)
        
        # 过滤掉已存在的数据
        df = df[~df['trade_date'].isin(existing_dates['trade_date'])]
        print('需要存储的新数据日期:')
        print(df['trade_date'])
        
        if not df.empty:
            # 将新数据追加到数据库
            df.to_sql('stock_index', con=engine, if_exists='append', index=False)
            print(f'成功存储{len(df)}条新数据！')
        else:
            print('没有新数据需要存储')
        print('数据存储成功！')
    except Exception as e:
        print(f'程序出错：{str(e)}')

if __name__ == '__main__':
    main()